from ultralytics.data.augment import RandomPerspective
from ultralytics.data import YOLODataset
from ultralytics.cfg import cfg2dict
import random
import cv2

# 初始化RandomPerspective类
transform = RandomPerspective(degrees=0.0, translate=0.0, scale=0.0, shear=0.0, perspective=0.005)  # 默认不做透视操作

# 输入数据labels
# 测试数据获取
config = cfg2dict('./coco8/coco8.yaml')
dataset = YOLODataset(data=config, img_path='./coco8/images')

# 随机选一张图片
idx = random.choices([0,1,2,3],k=1)[0]
idx = 2
labels = {"img":dataset.get_image_and_label(idx)['img'],
          "cls":dataset.get_image_and_label(idx)['cls'],
          "instances":dataset.get_image_and_label(idx)['instances']}

# 变化操作
results = transform(labels)
# 可视化保存
savedir = './randomperspective_test.jpg'
cv2.imwrite(savedir, results['img'])

# 框标签可视化保存
def vis_box(box_list, img, color):
    for box in box_list:
        l,t,r,b = box
        cv2.rectangle(img, (int(l), int(t)), (int(r), int(b)), color, thickness=2, lineType=cv2.LINE_AA)
    return img

# 修正框可视化
box_list = results['instances']._bboxes.bboxes
visdir = './randomperspective_vis.jpg'
img = cv2.imread(savedir)
cv2.imwrite(visdir, vis_box(box_list,img, (0,255,0)))